# -*- coding: utf-8 -*-

import tensorflow as tf
import numpy as np
from scipy.special import expit
from scipy import stats
import cv2
from vgg import vgg_network


def read_image(image_path):
    img = cv2.imread(image_path)
    img = img[540:, 0:960, :]
    img = cv2.resize(img, (320, 180))
    img = img.astype(np.float32)
    img = (img - [104, 110, 109]) / 255
    assert isinstance(img, np.ndarray)
    img = img.reshape([1] + list(img.shape))
    return img


def show_predict(img_path, predict):
    predict = expit(predict)
    predict = stats.threshold(predict, 0.6, None, 0)
    min_value, max_value = np.min(predict), np.max(predict)
    print(min_value, max_value)
    img = ((predict - min_value) / max_value) * 255
    img = img.astype(np.uint8, copy=True)

    src_img = cv2.imread(img_path)
    src_img = src_img[540:, 0:960, :]
    src_img = cv2.resize(src_img, (320, 180))

    show_image = np.zeros(shape=(180, 320, 3), dtype=np.uint8)
    show_image = show_image + [0, 255, 255]
    show_image = show_image.astype(dtype=np.uint8, copy=False)
    tmp = np.zeros_like(show_image)
    np.copyto(tmp, show_image, where=img > 0)

    cv2.addWeighted(tmp, 0.2, src_img, 0.8, 0, src_img)

    # cv2.imshow("show", src_img)
    # cv2.waitKey()
    return src_img


def main():
    test_x_input = tf.placeholder(dtype=tf.float32, shape=[None, 180, 320, 3],
                                  name="test_x_input")
    vgg_net, predict_node = vgg_network(test_x_input, batch_size=1)
    saver = tf.train.Saver()
    # cv2.namedWindow("show", cv2.WINDOW_NORMAL)
    with tf.Session() as sess:
        for i in range(1, 50):
            saver.restore(sess, "./model/vehicle_L2Loss.npk-981")
            # image_path = "/home/lijun/Dataset/vehicle_segment/origin/{:>07}.jpg".format(i + 1)
            image_path = "/home/lijun/Dataset/AK1126/{:>07}.jpg".format(i + 1)
            img = read_image(image_path)
            predict = sess.run(predict_node, feed_dict={test_x_input: img})
            predict = predict[0]
            show_img = show_predict(image_path, predict)
            save_path = "./test_results/{:>07d}.jpg".format(i)
            show_img = cv2.resize(show_img, (640, 360))
            cv2.imwrite(save_path, show_img)
    # cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
